#include "kernel_operator.h"
#include "lib/matmul_intf.h"

using namespace AscendC;

class NLLLossNew {
 public:
  	__aicore__ inline NLLLossNew() {}
  	__aicore__ inline void Init(GM_ADDR x, GM_ADDR target, GM_ADDR weight,
                              GM_ADDR y, GM_ADDR workspace, uint32_t reduction,
                              int32_t ignoreIndex, uint32_t sizeC,
                              uint32_t rowSmall, uint32_t rowBig, uint32_t rowMod, 
                              uint32_t coreNum) {

    
    this->reduction = reduction;
    this->ignoreIndex = ignoreIndex;
	this->sizeC = sizeC;
	this->rowSmall = rowSmall;
	this->rowBig = rowBig;
	this->rowMod = rowMod;
	this->coreNum = coreNum;

	this->coreIndex = GetBlockIdx();
	this->beginRow = (coreIndex < rowMod) ? (rowBig * this->coreIndex) : (rowSmall * this->coreIndex + rowMod);
	this->coreRow = (coreIndex < rowMod) ? rowBig : rowSmall;

    weightGm.SetGlobalBuffer((__gm__ DTYPE_WEIGHT*)weight,
                             sizeC * sizeof(DTYPE_WEIGHT));
    this->roundSizeC = ((this->sizeC+7) / 8) * 8;
    this->roundCoreRow = ((this->coreRow+7) / 8) * 8;
    this->roundCoreNum =  ((this->coreNum+7) / 8) * 8;
    // pipe.InitBuffer(AllWeightBuf, 2, roundSizeC * sizeof(DTYPE_WEIGHT));

	// this->allWeight =
	// AllWeightBuf.AllocTensor<DTYPE_WEIGHT>();
	// this->allWeight = AllWeightBuf.Get<DTYPE_WEIGHT>();
	// DataCopy(allWeight, weightGm[0], roundSizeC);

    targetGm.SetGlobalBuffer(
        (__gm__ DTYPE_TARGET*)target + this->beginRow,
        this->coreRow * sizeof(DTYPE_TARGET));
	
	// pipe.InitBuffer(AllTargetBuf,2, roundCoreRow * sizeof(DTYPE_TARGET));
	// pipe.InitBuffer(AllXBuf ,2, ((this->coreRow) / 8) * 8 * sizeof(DTYPE_TARGET));

// 	this->allTarget =
//         AllTargetBuf.AllocTensor<DTYPE_TARGET>();

	// this->allTarget = AllTargetBuf.Get<DTYPE_TARGET>();

	// pipe.InitBuffer(inQueueTarget, BUFFER_NUM, ITER_N * sizeof(DTYPE_TARGET));

	// pipe.InitBuffer(GatherWeightBuf, roundCoreRow * sizeof(DTYPE_WEIGHT));
	// this->gatherWeight = GatherWeightBuf.Get<DTYPE_WEIGHT>();

	// pipe.InitBuffer(InitIndexBuf, roundCoreRow * sizeof(DTYPE_TARGET));
	// this->initIndex = InitIndexBuf.Get<DTYPE_TARGET>();

	// pipe.InitBuffer(IndexBuf, this->coreRow * sizeof(DTYPE_TARGET));
	xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x + this->beginRow * this->sizeC,
                        this->coreRow * this->sizeC * sizeof(DTYPE_X));
	

	// pipe.InitBuffer(GatherXBuf, roundCoreRow * sizeof(DTYPE_X)*8);
	// this->gatherX = GatherXBuf.Get<DTYPE_X>();

	// pipe.InitBuffer(ReduceXBuf, 32);
	// pipe.InitBuffer(ReduceWBuf, 32);
	// pipe.InitBuffer(WorkBuf, ((this->coreNum+this->coreRow+7) / 8) * 8 * sizeof(DTYPE_X));

	// pipe.InitBuffer(FinalXBuf, this->roundCoreNum * sizeof(DTYPE_X));
	// pipe.InitBuffer(FinalWBuf, this->roundCoreNum * sizeof(DTYPE_X));

    yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y, sizeof(DTYPE_Y));

	// workXGm.SetGlobalBuffer((__gm__ DTYPE_X*)workspace, this->roundCoreNum * sizeof(DTYPE_X));
	// workWGm.SetGlobalBuffer(
	// (__gm__ DTYPE_WEIGHT*)workspace + this->roundCoreNum+ 32,
	// this->roundCoreNum * sizeof(DTYPE_WEIGHT));
	}

	__aicore__ inline void process() {
		// if (this->coreIndex >= 40) return;
		// printf("----info----\n coreNum=%d, coreIdx=%d, coreRow=%d \n",this->coreNum,this->coreIndex,this->coreRow);
		// copyInTarget();
		// GatherWeightAndX();
		// Compute();
		// CopyToWorkSpace();
		// SyncAll();
		// if (this->coreIndex != 0) return;
		// CopyFromWorkSpace();
		// ReduceFinal();
	}

	__aicore__ void copyInTarget() {
		// LocalTensor<DTYPE_TARGET> targetLocal = inQueueTarget.AllocTensor<DTYPE_TARGET>();
		DataCopy(this->allTarget, targetGm, (this->coreRow + 7) / 8 * 8);
		PipeBarrier<PIPE_ALL>();
		// printf("allTarget\n");
		// for (int i = 0; i < this->coreRow; i++) {
		// 	printf("%d ",allTarget.GetValue(i));
		// }
		// printf("\n");
	}

	__aicore__ void GatherWeightAndX() {
		// ------gather weight
		// CreateVecIndex(this->initIndex, 0, this->coreRow);
		if (0 <= this->ignoreIndex && this->ignoreIndex < this->sizeC) {
			this->allWeight.SetValue(this->ignoreIndex, 0);
			}

		auto indexLocal = InitIndexBuf.Get<DTYPE_TARGET>();
		Muls(indexLocal, this->allTarget, 4 , this->coreRow);
		auto indexLocalUint = indexLocal.ReinterpretCast<uint32_t>();
		PipeBarrier<PIPE_ALL>();
		// printf("----index\n");
		// for (int i = 0; i < this->coreRow; i++) {
		// 	printf("%d ",indexLocalUint.GetValue(i));
		// }
		// printf("\n");
		// PipeBarrier<PIPE_ALL>();
		// printf("size: %d %d %d\n", this->gatherWeight.GetSize(), this->allWeight.GetSize(),indexLocalUint.GetSize());
		Gather(this->gatherWeight, this->allWeight, indexLocalUint, (uint32_t)0, this->coreRow);
		// for (int i = 0; i < this->coreRow; i++) {
		// 	printf("i: %d %f \n",indexLocalUint.GetValue(i), gatherWeight.GetValue(i));
		// }
		// ------gather X
        AscendC::DataCopyPadParams padParams{false, 0, 0, 0};
        AscendC::DataCopyParams copyParams{1, 4, 0, 0};

        for (int i = 0; i < coreRow; i++) {
            AscendC::DataCopyPad(this->gatherX[i*8], xGm[i * this->sizeC + this->allTarget.GetValue(i)], copyParams, padParams);
        }
		Gather(this->gatherWeight, this->allWeight, indexLocalUint, (uint32_t)0, this->coreRow);

		// for(int i=0;i<this->coreRow;i++){
		// 	this->gatherX.SetValue(i, xGm.GetValue(i * this->sizeC + this->allTarget.GetValue(i)));
		// }
		PipeBarrier<PIPE_ALL>();
	}

	__aicore__ void Compute() {
		auto workLocal = WorkBuf.Get<DTYPE_X>();
		auto reduceXLocal = ReduceXBuf.Get<DTYPE_X>();
		auto reduceWLocal = ReduceWBuf.Get<DTYPE_WEIGHT>();
		if (this->reduction == 2) {	//mean
			ReduceSum(reduceWLocal, this->gatherWeight, workLocal, this->coreRow);
		}

		Mul(this->gatherX, this->gatherX, this->gatherWeight, this->coreRow);
		// for (int i = 0; i < this->coreRow; i++) {
		// 	printf("i: %d %f %f\n",i, gatherX.GetValue(i),gatherWeight.GetValue(i));
		// }
		ReduceSum(reduceXLocal, this->gatherX, workLocal, this->coreRow);
    	// printf("reducexLocal: %f \n",reduceXLocal.GetValue(0));

		PipeBarrier<PIPE_ALL>();
	}

	__aicore__ void CopyToWorkSpace() {
		auto reduceXLocal = ReduceXBuf.Get<DTYPE_X>();
		auto reduceWLocal = ReduceWBuf.Get<DTYPE_WEIGHT>();
	    AscendC::DataCopyExtParams copyParams{1, 4, 0, 0, 0};
   		 AscendC::DataCopyPad(workXGm[this->coreIndex], reduceXLocal, copyParams);
		if (this->reduction == 2)  // mean
		{
		AscendC::DataCopyPad(workWGm[this->coreIndex], reduceWLocal, copyParams);
		// DataCopy(workWGm[this->coreIndex * ITER_N], LocaloutW, ITER_N);
		PipeBarrier<PIPE_ALL>();
		// printf("\n-------workWGm:\n");
		// for (int i = 0; i < this->IterNum; i++) {
		//   printf("%f ", workWGm.GetValue(this->coreIndex + i));
		// }
		}
	}

	__aicore__ void CopyFromWorkSpace() {
		auto finalXLocal = FinalXBuf.Get<DTYPE_X>();
		auto finalWLocal = FinalWBuf.Get<DTYPE_WEIGHT>();
		// workXGm.GetValue(0);
		DataCopy(finalXLocal, workXGm, (this->coreNum + 7)/ 8* 8);
		if (this->reduction == 2) {	//mean
			DataCopy(finalWLocal, workWGm, (this->coreNum + 7)/ 8* 8);
		}
		// for (int i = 0; i < this->coreRow; i++)
		// 	printf("final %f %f\n", finalXLocal.GetValue(i),finalWLocal.GetValue(i));
		PipeBarrier<PIPE_ALL>();

		// PipeBarrier<PIPE_ALL>();
	}

	__aicore__ void ReduceFinal() {
		auto finalXLocal = FinalXBuf.Get<DTYPE_X>();
		auto finalWLocal = FinalWBuf.Get<DTYPE_WEIGHT>();
		auto reduceXLocal = ReduceXBuf.Get<DTYPE_X>();
		auto reduceWLocal = ReduceWBuf.Get<DTYPE_WEIGHT>();
		auto workLocal = WorkBuf.Get<DTYPE_X>();
		float p = -1;
		if (this->reduction == 2) {	//mean
			ReduceSum(reduceWLocal, finalWLocal, workLocal, this->coreNum);
			p = p / reduceWLocal.GetValue(0);
		}

		ReduceSum(reduceXLocal, finalXLocal, workLocal, this->coreNum);
		// DataCopy(yGm, workWGm, (this->coreNum + 7)/ 8* 8);

		p = reduceXLocal.GetValue(0) * p;
    	yGm.SetValue(0, p);

	}


private:
	uint32_t reduction, sizeC;
	int32_t ignoreIndex;
	uint32_t rowSmall, rowBig, rowMod, coreNum,roundSizeC,roundCoreRow,roundCoreNum;
	uint32_t beginRow,coreIndex,coreRow;

	AscendC::TPipe pipe;
	AscendC::GlobalTensor<DTYPE_X> xGm;
	AscendC::GlobalTensor<DTYPE_TARGET> targetGm;
	AscendC::GlobalTensor<DTYPE_WEIGHT> weightGm;
	AscendC::GlobalTensor<DTYPE_Y> yGm, workXGm, workWGm;

	TQue<QuePosition::VECIN, 2> inQueueX, inQueueTarget, AllTargetBuf, AllWeightBuf;

	TBuf<QuePosition::VECCALC> InitIndexBuf,IndexBuf,GatherWeightBuf,GatherXBuf;
	TBuf<QuePosition::VECCALC> ReduceXBuf,ReduceWBuf,WorkBuf;
	TBuf<QuePosition::VECCALC> FinalXBuf,FinalWBuf;

	LocalTensor<DTYPE_WEIGHT> allWeight,gatherWeight;
	LocalTensor<DTYPE_TARGET> allTarget;
	LocalTensor<DTYPE_TARGET> initIndex;
	LocalTensor<DTYPE_X> gatherX;
	
};